package com.superai.controller;

import io.github.imfangs.dify.client.DifyCompletionClient;
import io.github.imfangs.dify.client.callback.CompletionStreamCallback;
import io.github.imfangs.dify.client.enums.ResponseMode;
import io.github.imfangs.dify.client.event.*;
import io.github.imfangs.dify.client.model.chat.AppInfoResponse;
import io.github.imfangs.dify.client.model.chat.AppParametersResponse;
import io.github.imfangs.dify.client.model.common.SimpleResponse;
import io.github.imfangs.dify.client.model.completion.CompletionRequest;
import io.github.imfangs.dify.client.model.completion.CompletionResponse;
import lombok.RequiredArgsConstructor;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 文本生成应用控制器
 *
 * 这个控制器提供了与Dify文本生成应用交互的API，包括：
 * - 发送文本生成请求（阻塞模式和流式模式）
 * - 获取应用信息和参数
 * - 停止文本生成
 */
@RestController
@RequestMapping("/completion")
@RequiredArgsConstructor
public class CompletionController {

    /**
     * Dify文本生成客户端，通过构造函数注入
     */
    @Autowired
    private  DifyCompletionClient completionClient;

    /**
     * 存储SSE连接的映射表，用于流式响应
     */
    private final ConcurrentHashMap<String, SseEmitter> sseEmitters = new ConcurrentHashMap<>();

    /**
     * 发送文本生成请求（阻塞模式）
     *
     * @param request 文本生成请求
     * @return 文本生成响应
     */
    @PostMapping
    public ResponseEntity<CompletionResponse> sendCompletionMessage(@RequestBody CompletionRequest request) {
        try {
            // 设置响应模式为阻塞
            request.setResponseMode(ResponseMode.BLOCKING);

            // 发送请求并获取响应
            CompletionResponse response = completionClient.sendCompletionMessage(request);
            return ResponseEntity.ok(response);
        } catch (IOException e) {
            // 记录异常并返回错误响应
            return ResponseEntity.internalServerError().build();
        }
    }

    /**
     * 发送文本生成请求（流式模式）
     *
     * @param request 文本生成请求
     * @return SSE发射器，用于流式传输响应
     */
    @PostMapping(value = "/stream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public SseEmitter sendCompletionMessageStream(@RequestBody CompletionRequest request) {
        // 创建SSE发射器，设置超时时间为30分钟
        SseEmitter emitter = new SseEmitter(1800000L);
        String emitterId = request.getUser() + "_" + System.currentTimeMillis();
        sseEmitters.put(emitterId, emitter);

        try {
            // 设置响应模式为流式
            request.setResponseMode(ResponseMode.STREAMING);

            // 发送流式请求
            completionClient.sendCompletionMessageStream(request, new CompletionStreamCallback() {
                @Override
                public void onMessage(MessageEvent event) {
                    try {
                        emitter.send(SseEmitter.event()
                                .name("message")
                                .data(event));
                    } catch (IOException e) {
                        onException(e);
                    }
                }

                @Override
                public void onMessageEnd(MessageEndEvent event) {
                    try {
                        emitter.send(SseEmitter.event()
                                .name("message_end")
                                .data(event));
                        emitter.complete();
                    } catch (IOException e) {
                        onException(e);
                    } finally {
                        sseEmitters.remove(emitterId);
                    }
                }

                @Override
                public void onTtsMessage(TtsMessageEvent event) {
                    try {
                        emitter.send(SseEmitter.event()
                                .name("tts_message")
                                .data(event));
                    } catch (IOException e) {
                        onException(e);
                    }
                }

                @Override
                public void onTtsMessageEnd(TtsMessageEndEvent event) {
                    try {
                        emitter.send(SseEmitter.event()
                                .name("tts_message_end")
                                .data(event));
                    } catch (IOException e) {
                        onException(e);
                    }
                }

                @Override
                public void onMessageReplace(MessageReplaceEvent event) {
                    try {
                        emitter.send(SseEmitter.event()
                                .name("message_replace")
                                .data(event));
                    } catch (IOException e) {
                        onException(e);
                    }
                }

                @Override
                public void onError(ErrorEvent event) {
                    try {
                        emitter.send(SseEmitter.event()
                                .name("error")
                                .data(event));
                        emitter.complete();
                    } catch (IOException e) {
                        onException(e);
                    } finally {
                        sseEmitters.remove(emitterId);
                    }
                }

                @Override
                public void onException(Throwable throwable) {
                    emitter.completeWithError(throwable);
                    sseEmitters.remove(emitterId);
                }

                @Override
                public void onPing(PingEvent event) {
                    try {
                        emitter.send(SseEmitter.event()
                                .name("ping")
                                .data(event));
                    } catch (IOException e) {
                        // 忽略心跳事件的异常
                    }
                }
            });
        } catch (IOException e) {
            emitter.completeWithError(e);
            sseEmitters.remove(emitterId);
        }

        // 设置超时和完成回调
        emitter.onTimeout(() -> sseEmitters.remove(emitterId));
        emitter.onCompletion(() -> sseEmitters.remove(emitterId));

        return emitter;
    }

    /**
     * 停止文本生成
     *
     * @param taskId 任务ID
     * @param userId 用户ID
     * @return 简单响应
     */
    @PostMapping("/stop")
    public ResponseEntity<SimpleResponse> stopCompletion(
            @RequestParam String taskId,
            @RequestParam String userId) {
        try {
            SimpleResponse response = completionClient.stopCompletion(taskId, userId);
            return ResponseEntity.ok(response);
        } catch (IOException e) {
            return ResponseEntity.internalServerError().build();
        }
    }

    /**
     * 获取应用信息
     *
     * @return 应用信息响应
     */
    @GetMapping("/app-info")
    public ResponseEntity<AppInfoResponse> getAppInfo() {
        try {
            AppInfoResponse response = completionClient.getAppInfo();
            return ResponseEntity.ok(response);
        } catch (IOException e) {
            return ResponseEntity.internalServerError().build();
        }
    }

    /**
     * 获取应用参数
     *
     * @return 应用参数响应
     */
    @GetMapping("/app-parameters")
    public ResponseEntity<AppParametersResponse> getAppParameters() {
        try {
            AppParametersResponse response = completionClient.getAppParameters();
            return ResponseEntity.ok(response);
        } catch (IOException e) {
            return ResponseEntity.internalServerError().build();
        }
    }

    /**
     * 文字转语音
     *
     * @param request 文字转语音请求
     * @return 音频数据
     */
    @PostMapping(value = "/text-to-audio", produces = MediaType.APPLICATION_OCTET_STREAM_VALUE)
    public ResponseEntity<byte[]> textToAudio(@RequestBody TextToAudioRequest request) {
        try {
            byte[] audioData = completionClient.textToAudio(
                    request.getConversationId(),
                    request.getText(),
                    request.getUserId());
            return ResponseEntity.ok()
                    .header("Content-Disposition", "attachment; filename=audio.mp3")
                    .body(audioData);
        } catch (IOException e) {
            return ResponseEntity.internalServerError().build();
        }
    }

    /**
     * 文字转语音请求类
     */
    public static class TextToAudioRequest {
        private String conversationId;
        private String text;
        private String userId;

        // Getters and setters
        public String getConversationId() {
            return conversationId;
        }

        public void setConversationId(String conversationId) {
            this.conversationId = conversationId;
        }

        public String getText() {
            return text;
        }

        public void setText(String text) {
            this.text = text;
        }

        public String getUserId() {
            return userId;
        }

        public void setUserId(String userId) {
            this.userId = userId;
        }
    }
}
